import torch.utils.data as data
import torchvision.transforms as transforms
from utils.config import *
from dataloader.KITTI2015_loader import KITTI2015
from dataloader.SceneFlow_loader import SceneFlow
from dataloader.USVInland_loader import USVInland
from dataloader.SGBM_loader import USVInland_SGBM

# 数据预处理
transform = transforms.Compose([transforms.ToTensor()]) # 组合操作：转换张量

# 准备数据集
if dataset == 'kitti2015':
    if system == 'Windows':
        root = 'G:/KITTI/KITTI2015/data_scene_flow' # Windows
    elif system == 'Ubuntu':
        root = '/home/ubuntu/zhouyiqing/KITTI/KITTI2015/data_scene_flow' # Ubuntu
    img_dataset = KITTI2015(root=root, transform=transform, occ=True)

elif dataset == 'sceneflow':
    if system == 'Windows':
        root = 'G:/SceneFlow' # Windows
    elif system == 'Ubuntu':
        root = '/home/ubuntu/zhouyiqing/SceneFlow' # Ubuntu
    img_dataset = SceneFlow(root=root, transform=transform)

elif dataset == 'usvinland':
    if system == 'Windows':
        root = 'G:/USVInland/Stereo Matching/Low_Res_640_320' # Windows
    elif system == 'Ubuntu':
        root = '/home/ubuntu/zhouyiqing/USVInland/Stereo Matching/Low_Res_640_320' # Ubuntu
    img_dataset = USVInland(root=root, transform=transform, seg=False)

elif dataset == 'usvinland_seg':
    if system == 'Windows':
        root = 'G:/USVInland/Stereo Matching/Segmentation' # Windows
    elif system == 'Ubuntu':
        root = '/home/ubuntu/zhouyiqing/USVInland/Stereo Matching/Segmentation' # Ubuntu
    img_dataset = USVInland(root=root, transform=transform, seg=True)

if dataset == 'sgbm':
    assert system == 'Windows', 'system must be Windows'
    sgbm_root = 'G:/USVInland/Stereo Matching/dispVisualization/SGBM_Disp_Map' # Windows
    gt_root = 'G:/USVInland/Stereo Matching/Low_Res_640_320/Disp_Map'
    img_dataset = USVInland_SGBM(sgbm_root=sgbm_root, gt_root=gt_root)

# 数据加载器
img_loader = data.DataLoader(img_dataset, batch_size=1, shuffle=True, num_workers=0, drop_last=True)
